# -*- coding: utf-8 -*-
# @Time    : 2025/1/8 下午2:42
# @Author  : ysj
# @FileName: preson_defect.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/ydscc?type=blog
# yolov5
# 导包
import torch
import cv2
from multiprocessing import Process, Manager, Value
# 下面两个是yolov5文件夹里面的代码
from utils.general import non_max_suppression
from models.experimental import attempt_load
import numpy as np
import logging

import platform
import pathlib
plt = platform.system()
if plt == 'Windows':
   pathlib.PosixPath = pathlib.WindowsPath



def boxes_intersect(box1, box2):
    """
    判断两个矩形框是否相交
    :param box1: [x1, y1, x2, y2] 左上角和右下角坐标
    :param box2: (x, y, width, height) 左上角坐标和宽高
    :return: 是否相交 (True/False)
    """
    # 解析 box1
    x1_min, y1_min, x1_max, y1_max = box1
    # 提取 box2 的坐标
    x2_min, y2_min = box2[0], box2[1]
    x2_max, y2_max = x2_min + box2[2], y2_min + box2[3]

    # 检查是否有交集
    horizontal_overlap = not (x1_max < x2_min or x2_max < x1_min)
    vertical_overlap = not (y1_max < y2_min or y2_max < y1_min)

    return horizontal_overlap and vertical_overlap


def calculate_min_distance(box1, box2):
    """
    计算两个矩形框之间的最近距离
    :param box1: 矩形框1，格式 [x1, y1, x2, y2]
    :param box2: 矩形框2，格式 [x1, y1, x2, y2]
    :return: 最近距离
    """
    # 矩形框的左上角和右下角
    x1_min, y1_min, x1_max, y1_max = box1
    x2_min, y2_min, x2_max, y2_max = box2

    # 判断是否有重叠
    if x1_min < x2_max and x1_max > x2_min and y1_min < y2_max and y1_max > y2_min:
        return 0  # 如果有重叠，距离为 0

    # 分别计算在 x 和 y 方向上的最小间距
    dx = max(0, max(x2_min - x1_max, x1_min - x2_max))  # x方向
    dy = max(0, max(y2_min - y1_max, y1_min - y2_max))  # y方向

    # 返回欧几里得距离
    return np.sqrt(dx**2 + dy**2)

# 确保在进行对象检测时，边界框的位置可以与输入图像的大小相对应，以便正确地识别和定位对象。
def scale_coords(img, coords, out_shape):
    # img: 输入图像，具有(C, H, W)的形状，表示通道数、高度和宽度。
    # coords: 包含边界框坐标的数组，具有(..., 4)的形状，其中...表示可变数量的维度，而最后一个维度4表示左上角和右下角的坐标。
    # out_shape: 输出尺寸，一个包含高度、宽度和通道数的元组或列表([out_h, out_w, _])。
    img_h, img_w = img.shape[2:]
    out_h, out_w, _ = out_shape
    coords[..., 0] *= out_w / img_w
    coords[..., 1] *= out_h / img_h
    coords[..., 2] *= out_w / img_w
    coords[..., 3] *= out_h / img_h
    # 通过按比例缩放坐标值来调整边界框的位置。首先，将输入图像的高度和宽度记录为img_h和img_w。
    # 然后，获取输出尺寸的高度和宽度，分别记录为out_h和out_w。
    # 接下来，通过计算比例因子out_w / img_w和out_h / img_h，将边界框的所有坐标值乘以相应的比例因子。
    # 这样做的目的是将坐标值映射到新的输出尺寸上，以便适应不同尺寸的输入图像。
    return coords
    # 最后，返回经过缩放的坐标数组coords

def img_padd(image,target_size=(640,640)):

    """

    Args:
        image: 图像
        target_size: 等比压缩，填充尺寸

    Returns:

    """
    h, w = image.shape[:2]

    # 计算缩放比例和填充
    scale = min(target_size[0] / w, target_size[1] / h)
    new_w, new_h = int(w * scale), int(h * scale)
    pad_w, pad_h = target_size[0] - new_w, target_size[1] - new_h
    pad_left, pad_top = pad_w // 2, pad_h // 2

    # 缩放图像
    resized_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    # 填充图像
    padded_image = cv2.copyMakeBorder(
        resized_image, pad_top, pad_h - pad_top, pad_left, pad_w - pad_left,
        cv2.BORDER_CONSTANT, value=(114, 114, 114)  # 灰色填充
    )
    return padded_image


def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, stride=32):
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))

    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding

    dw /= 2  # divide padding into 2 sides
    dh /= 2
    print(dw, dh)
    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)

def scale_bboxes(bboxes, scale, padding, original_size=(1920, 1080)):
    left, top = padding
    target_w, target_h = 640, 640

    # 反向操作：去除填充的影响，按照缩放比例还原回原图尺寸


    x1, y1, x2, y2 = bboxes
    # 去除填充
    x1 = (x1 - left) / scale
    y1 = (y1 - top) / scale
    x2 = (x2 - left) / scale
    y2 = (y2 - top) / scale

    # 确保边界框不超出原图大小
    x1 = max(0, x1)
    y1 = max(0, y1)
    x2 = min(original_size[0], x2)
    y2 = min(original_size[1], y2)

    # 将边界框添加到列表
    bboxes_scaled=[x1, y1, x2, y2]
    bboxes_scaled=[int(x) for x in bboxes_scaled]

    return bboxes_scaled

def detect_objects(weights_path, output_frame):

    import  time

    logging.basicConfig(level=logging.INFO)

    boxroi = (1000, 400, 440, 280)
    # x, y, w, h = 1000, 400, 440, 280
    text = "On duty"
    text1 = "Leave  Post"

    # 文本位置 (x, y)
    font = cv2.FONT_HERSHEY_SIMPLEX  # 字体类型
    font_scale = 3  # 字体大小
    color = (255, 255, 0)  # 字体颜色 (BGR)，此处为绿色
    thickness = 5  # 字体线宽
    # 加载YOLOv5模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 检查是否可用CUDA加速。如果CUDA可用 ，则将设备设置为'cuda'，否则设置为'cpu'。
    model = attempt_load(weights_path, device)

    # 设置置信度阈值和IoU阈值
    model.conf = 0.2
    model.iou = 0.6

    # 设置模型为推理模式
    model.eval()
    rtsp_url = ("rtsp://admin:ascendit123@10.10.70.28:554/h264/ch1/main/av_stream")


    cap = cv2.VideoCapture(rtsp_url)

    # 打开视频流
    # cap = cv2.VideoCapture(r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\runs\output_video1.avi')
    cap.set(3, 640)  # 设置分辨率
    cap.set(4, 640)

    while True:
        s = time.time()
        ret, frame = cap.read()  # 视频读入

        if not ret:
            break
        # cv2.rectangle(frame, (boxroi[0], boxroi[1]), (boxroi[0]+boxroi[3] , boxroi[1] + boxroi[2]), (0, 255, 0), 2)  # 绘制一个绿色矩形框

        # 图像预处理
        # frame = img_padd(frame)
        img = frame.copy()
        # img = img_padd(img)

        # 创建一个副本以进行预处理
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = torch.from_numpy(img).to(device).float() / 255.0
        img = img.permute(2, 0, 1).unsqueeze(0)
        img = torch.nn.functional.interpolate(img, size=(384, 640), mode='bilinear', align_corners=False)

        #输入大小为等比放缩之后的大小 且要被32 整除
        # size = (320, 320)：表示目标输出图像的大小为320x320像素。将输入图像按照此大小进行调整。
        # mode = 'bilinear'：表示插值方法采用双线性插值。双线性插值是一种常用的插值方法，它通过对最近的四个像素进行加权平均来计算新像素值。
        # align_corners = False：表示在计算插值时不对齐角点。这个参数的具体影响取决于插值方法和具体的实现方式。
        # 进行物体检测
        outputs = model(img)
        results = non_max_suppression(outputs, conf_thres=model.conf, iou_thres=model.iou)
        # 通过非最大抑制算法对模型输出的边界框进行处理，以选择最佳的边界框并过滤掉冗余的检测结果。
        # outputs：模型的输出结果，通常是一个包含预测边界框及其置信度的数组。
        # conf_thres：表示置信度阈值，用于过滤掉置信度低于此阈值的边界框。
        # iou_thres：表示IoU（Intersection over Union）阈值，用于合并重叠度高于此阈值的边界框。
        # 算法的步骤如下：
        # 1、根据置信度阈值过滤掉置信度低于阈值的边界框。
        # 2、对剩余的边界框按照置信度进行排序，置信度高的排在前面。
        # 3、从排好序的边界框列表中选择置信度最高的边界框，并将其添加到最终的结果列表中。
        # 4、遍历剩余的边界框列表，计算当前边界框与已选择的边界框的IoU值（重叠度），如果大于IoU阈值，则将其从列表中移除。
        # 5、重复步骤3和4，直到所有边界框都被处理完毕。
        # 6、最终，non_max_suppression()函数返回经过非最大抑制处理后的边界框结果。
        if results[0] is not None and len(results[0]) > 0:
            result = results[0]
            # 绘制边界框和标
        dis_res = 0
        ori_xyxy = None
        xyxy_list = []
        dis_list = []
        label_id =0
        for result in results:  # 逐个访问并获取
            '''
            所有检测矩形框进行排序，判断哪个离roi_id 最近，在进行判断相交的问题
            最后在画出最终结果 
            '''
            if result is not None and len(result) > 0:
                result[:, :4] = scale_coords(img, result[:, :4], frame.shape).round()
                # result[:, :4]表示选取每个边界框的前四个元素，即边界框的坐标信息。
                for *xyxy, conf, cls in result:
                        label_id+=1
                        # result是一个数组或张量，每一行代表一个边界框，包含边界框的坐标、类别标签和置信度等信息。
                        # *xyxy表示将边界框的前四个元素解包为一个名为xyxy的列表。这里的xyxy表示边界框的坐标信息，通常是左上角和右下角的坐标值。
                        # conf表示边界框的置信度，通常是模型对该边界框所属类别的预测置信度。
                        # cls表示边界框的类别标签，通常是模型对该边界框所属类别的预测结果。
                        label = f'{model.names[int(cls)]} {conf:.2f}'
                        label = label+f"_{label_id}"
                        # f'{model.names[int(cls)]} {conf:.2f}'使用了格式化字符串的语法，将类别名称和置信度转换为一个字符串。
                        # 其中，{model.names[int(cls)]}表示插入类别名称，
                        # {conf: .2f}表示插入置信度，并保留两位小数。
                        # 绘制边界框
                        xyxy = [int(x) for x in xyxy]
                        # xyxy = scale_bboxes(xyxy, scale, padding, original_size=(1920, 1080))
                        cv2.rectangle(frame, (xyxy[0], xyxy[1]), (xyxy[2], xyxy[3]), (0, 255, 0), 2)
                        cv2.putText(frame, label, (xyxy[0], xyxy[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                                    (0, 255, 0), 2)





        cv2.imshow('RTSP Video Stream with Person Detection1', frame)
        cv2.imwrite("output.jpg", frame)
            # cv2.imwrite('output.png', frame)
            # break
        logging.info(f"用时：{time.time()-s:.2f}")
        cv2.waitKey(1)
        # 按q退出
        if cv2.waitKey(1) == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # 设置权重文件路径
    weights_path = r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\runs\best_person_stuta.pt'
    #weights_path = r'C:\Al\Software\AI_Model\Project\yolov5\yolov5\best.pt'
    # 使用 Manager 创建共享变量
    manager = Manager()
    output_frame = manager.dict()

    # 创建物体检测进程
    detection_process = Process(target=detect_objects, args=(weights_path, output_frame))
    detection_process.start()
    # 等待物体检测进程结束
    detection_process.join()

